package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"strconv"
	"strings"
	"time"
)

/* 构建第一层代理 */

func main() {
	addr := "http://127.0.0.1:2008"
	rs1 := []string{"http://127.0.0.1:2003", "http://127.0.0.1:2004"}
	urls := []*url.URL{}
	for _, i := range rs1 {
		url1, err1 := url.Parse(i)

		if err1 != nil {
			fmt.Println(err1)
			continue
		}

		urls = append(urls, url1)
	}

	proxy := NewMultipleHostReverseProxy(urls)
	fmt.Println("Starting httpserver at " + addr)
	fmt.Println(http.ListenAndServe(":2001", proxy))

}

/* 这里就是真实构建完成反向代理服务 */
/* 步骤：

 */

// 构建连接池对象
var transport = &http.Transport{
	DialContext: (&net.Dialer{
		Timeout:   30 * time.Second, //连接超时
		KeepAlive: 30 * time.Second, //长连接超时时间
	}).DialContext,
	MaxIdleConns:          100,              //最大空闲连接
	IdleConnTimeout:       90 * time.Second, //空闲超时时间
	TLSHandshakeTimeout:   10 * time.Second, //tls握手超时时间
	ExpectContinueTimeout: 1 * time.Second,  //100-continue 超时时间
}

// 构建构建RP对象
func NewMultipleHostReverseProxy(targets []*url.URL) *httputil.ReverseProxy {
	// 填充director /* 对请求属性的修改 ， 处理到这一层，就直接修改掉req中的部分结构内容*/
	director := func(req *http.Request) {
		// 随机负载均衡
		ind := rand.Intn(len(targets))
		target := targets[ind]

		// 基本属性填充(覆盖)
		req.URL.Scheme = target.Scheme
		req.URL.Host = target.Host

		// Url重写——毕竟做的是代理
		req.URL.Path = JoinUrlQuery(target.Path, req.URL.Path) // 这里注意顺序的情况——并且是针对path的操作

		// 针对Query的处理
		// 如果原URL的query是"",代理的URL也是"",则直接拼接
		// 如果不为空，需要增加 & 符号拼接
		targetQuery := target.RawQuery
		if targetQuery == "" || req.URL.RawQuery == "" {
			req.URL.RawQuery = targetQuery + req.URL.RawQuery
		} else {
			req.URL.RawQuery = targetQuery + "&" + req.URL.RawQuery
		}

		// 设置User-Agent
		if _, ok := req.Header["User-Agent"]; !ok {
			req.Header.Set("User-Agent", "user-agent")
		}

		// 第一代理需要设置x-real-ip TODO
		//只在第一代理中设置此header头
		req.Header.Set("X-Real-Ip", req.RemoteAddr)

	}

	// 更改请求返回的内容
	// 针对正常返回的内容，只需要传递就可以——至于返回中的IP等元数据填充的，应该是在拿到返回具体返回内容之后填充
	// 这里modify不需要管
	modifyFunc := func(resp *http.Response) error {
		var payload []byte
		var readErr error
		payload, readErr = ioutil.ReadAll(resp.Body)

		// 不是200请求的时候
		if resp.StatusCode != 200 {
			//获取内容
			payload = []byte("StatusCode error:" + string(payload))
		}

		//追加内容

		// 预读了数据，所以需要内容重新回写Header
		resp.Body = ioutil.NopCloser(bytes.NewBuffer(payload))
		resp.ContentLength = int64(len(payload))
		resp.Header.Set("Content-Length", strconv.FormatInt(int64(len(payload)), 10))

		return readErr
	}

	//错误回调 ：关闭real_server时测试，错误回调
	errFunc := func(w http.ResponseWriter, r *http.Request, err error) {
		http.Error(w, "ErrorHandler error:"+err.Error(), 500)
	}

	return &httputil.ReverseProxy{
		Director:       director,
		Transport:      transport,
		ModifyResponse: modifyFunc,
		ErrorHandler:   errFunc,
	}
}

func JoinUrlQuery(a, b string) string {
	al := strings.HasSuffix(a, "/") // 检查尾部
	bl := strings.HasPrefix(b, "/") // 检查头部
	if al && bl {                   // /aaa/  /bbb 的情况
		return a + b[1:]
	}
	if !al && !bl { // /aaa  bbb/ 的情况
		return a + "/" + b
	}
	return a + b // /aaa、/bbb  ||  /aaa/、bbb/
}
